#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

class KernelTril{
public:
    __aicore__ inline KernelTril() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, uint32_t xShape, uint32_t yShape,
                                int32_t value, uint32_t blockLength)
    {
        this->value = value;
        this->xShape = xShape;
        this->yShape = yShape;
	this->tileNum = tileNum;
        this->lasttileLength = lasttileLength;
        this->tileLength = tileLength / BUFFER_NUM;
        this->blockLength = blockLength;
        xGm.SetGlobalBuffer((__gm__ DTYPE_X *)x, blockLength);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)x, blockLength);
    	pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
    }

    __aicore__ inline void Process()
    {    
    int32_t loopCount = this->tileNum * BUFFER_NUM;
    for (int32_t i = 0; i < loopCount; i++) {
      CopyIn(i);
      Compute(i);
      CopyOut(i);
    }
    }

private:
    __aicore__ inline CopyIn(int32_t progress){
    LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
          (progress == (this->tileNum * BUFFER_NUM - 1))) {
        DataCopy(
            xLocal,
            xGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            (this->tileLength));
      }else
      {
	DataCopy(xLocal, xGm[progress * (this->tileLength)],
                (this->tileLength));
      }

	inQueueX.EnQue(xLocal);

    }

    __aicore__ inline void Compute(int32_t progress)
    {
	LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
        int32_t dim_xy = this->xShape * this->yShape;
        int32_t dim_f = this->blockLength / dim_xy;

        //for (int32_t f = 0; f < dim_f; f++)
        //{
         //   for (int32_t i = 0; i < this->xShape; i++)
          //  {
           //     for (int32_t j = i + this->value;j<this->yShape; j++)
            //    {
             //       yGm.SetValue(f * dim_xy + i * this->yShape + j, 0);
              //  }
           // }
       // }
    }

private:
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_Y> yGm;
    int32_t blockLength;
    int32_t tileNum;
    int32_t tileLength;
    int32_t lasttileLength;
    int32_t value;
    uint32_t xShape;
    uint32_t yShape;
};
extern "C" __global__ __aicore__ void tril(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tilingData, tiling);
    // TODO: user kernel impl
    KernelTril op;
    op.Init(x, y, tilingData.xShape, tilingData.yShape,
            tilingData.value, tilingData.blockLength);
    op.Process();
}
